# -*- coding: utf-8 -*-
"""
文件名称: sub_experiment/exp2/data_collect.py

功能概述:
    本脚本用于第一轮正式训练的数据准备, 主要完成以下工作:

    1. 读取两个已完成过滤的 Alpaca 格式数据集 (均为顶层 list[dict]):
        - med_qa_path: 公共数据集过滤后的 Simple QA 数据
        - aier_qa_path: Aier 数据过滤后的 Simple QA 数据

       样例结构 (单条记录):
        {
            "question_type": "compare",
            "instruction": "<image>......",
            "input": "",
            "output": "......",
            "images": [
                "/path/to/image.png"
            ]
        }

    2. 将两个数据集合并为一个总的数据集, 写入目标目录 out_dir 下的:
        - stage1_simple_qa.json

    3. 在同一目录下生成 LLaMA Factory 所需的 dataset_info.json,
       其内容固定为:
        {
          "stage1_simple_qa": {
            "file_name": "stage1_simple_qa.json",
            "columns": {
              "prompt": "instruction",
              "query": "input",
              "response": "output",
              "images": "images"
            }
          }
        }

    4. 进行简单统计:
        - 总样本数
        - 各来源数据集样本数 (med_qa / aier_qa)
        - question_type 分布及占比
          * question_type 字段不做先验假设, 遇到新的类型自动加入统计
          * 最终按频数从高到低排序

       统计结果以 JSON 格式写入:
        - out_dir / stage1_simple_qa_stats.json

       同时也会以 JSON 字符串形式打印到标准输出, 方便快速查看。

使用示例:
    python sub_experiment/exp2/data_collect.py \
        --med-qa-path /path/to/med_pub_simpleQA_filtered.json \
        --aier-qa-path /path/to/aier_simpleQA_filtered.json \
        -o /path/to/output_dir

依赖:
    - 标准库: os, json, argparse, logging, pathlib, typing
"""

import argparse
import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Tuple


logger = logging.getLogger(__name__)
if not logger.handlers:
    _handler = logging.StreamHandler()
    _formatter = logging.Formatter(
        fmt="[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    _handler.setFormatter(_formatter)
    logger.addHandler(_handler)
    logger.setLevel(logging.INFO)


# -----------------------------
# 基础工具函数
# -----------------------------

def _load_alpaca_json(path: Path, source_name: str) -> List[Dict[str, Any]]:
    """
    读取 Alpaca 格式 JSON 文件, 要求顶层结构为 list[dict]。

    参数:
        path:        JSON 文件路径
        source_name: 该数据源的名称 (仅用于错误信息与日志, 如 "med_qa" / "aier_qa")

    返回:
        records: list[dict] 形式的数据

    异常:
        - 文件不存在 / 读取失败 / JSON 解析失败 / 顶层非 list 时抛出异常
    """
    if not path.is_file():
        raise FileNotFoundError(f"{source_name} JSON 文件不存在: {path}")

    try:
        text = path.read_text(encoding="utf-8")
        data = json.loads(text)
    except Exception as e:
        raise RuntimeError(
            f"读取或解析 {source_name} JSON 失败: path={path}, error={e}"
        ) from e

    if not isinstance(data, list):
        raise RuntimeError(
            f"{source_name} JSON 顶层结构必须为 list, 实际为: {type(data)} (path={path})"
        )

    # 简单检查每一项是否为 dict
    for idx, item in enumerate(data, 1):
        if not isinstance(item, dict):
            raise RuntimeError(
                f"{source_name} JSON 索引 {idx} 的元素不是 dict 类型, 实际为: {type(item)}"
            )

    logger.info(
        "成功读取 %s 数据集: path=%s, 样本数=%d",
        source_name,
        path,
        len(data),
    )
    return data


def _build_dataset_info() -> Dict[str, Any]:
    """
    构造 LLaMA Factory 所需的 dataset_info.json 内容。

    返回:
        dataset_info: dict, 直接写入 JSON 即可
    """
    dataset_info = {
        "stage1_simple_qa": {
            "file_name": "stage1_simple_qa.json",
            "columns": {
                "prompt": "instruction",
                "query": "input",
                "response": "output",
                "images": "images",
            },
        }
    }
    return dataset_info


def _compute_stats(
    med_records: List[Dict[str, Any]],
    aier_records: List[Dict[str, Any]],
    combined_records: List[Dict[str, Any]],
) -> Dict[str, Any]:
    """
    统计数据规模与 question_type 分布。

    统计内容:
        - total_samples: 总样本数
        - by_source: { "med_qa": N1, "aier_qa": N2 }
        - question_type_distribution: list[{
              "question_type": str,
              "count": int,
              "ratio": float
          }]
          * question_type 动态收集, 遇到新的类型自动加入
          * 最终按 count 从高到低排序

    参数:
        med_records:     med_qa 数据列表
        aier_records:    aier_qa 数据列表
        combined_records:合并后的总列表

    返回:
        stats: 统计结果字典
    """
    total = len(combined_records)
    by_source = {
        "med_qa": len(med_records),
        "aier_qa": len(aier_records),
    }

    type_counter: Dict[str, int] = {}
    for idx, rec in enumerate(combined_records, 1):
        q_type = rec.get("question_type", None)
        # 不对 question_type 做强假设, 统一转为字符串; 若缺失则标记为 "__MISSING__"
        if q_type is None:
            q_type_key = "__MISSING__"
        else:
            q_type_key = str(q_type)

        type_counter[q_type_key] = type_counter.get(q_type_key, 0) + 1

    # 构造分布列表并按 count 从高到低排序
    dist_list: List[Dict[str, Any]] = []
    for q_type, cnt in type_counter.items():
        ratio = (cnt / float(total)) if total > 0 else 0.0
        dist_list.append(
            {
                "question_type": q_type,
                "count": int(cnt),
                "ratio": ratio,
            }
        )
    dist_list_sorted = sorted(dist_list, key=lambda x: x["count"], reverse=True)

    stats = {
        "total_samples": int(total),
        "by_source": by_source,
        "question_type_distribution": dist_list_sorted,
    }
    return stats


# -----------------------------
# CLI 解析与主流程
# -----------------------------

def parse_args() -> argparse.Namespace:
    """
    命令行参数解析。

    参数:
        --med-qa-path       公共数据集过滤后的 Simple QA JSON 路径 (顶层 list[dict])
        --aier-qa-path      Aier 数据过滤后的 Simple QA JSON 路径 (顶层 list[dict])
        --out-dir, -o       输出目录:
                               - stage1_simple_qa.json
                               - dataset_info.json
                               - stage1_simple_qa_stats.json
    """
    parser = argparse.ArgumentParser(
        description="Stage1 Simple QA 训练数据准备脚本: 合并两个 Alpaca 数据集并生成 dataset_info.json 与基本统计",
    )
    parser.add_argument(
        "--med-qa-path",
        type=str,
        required=True,
        help="公共数据集过滤后的 Simple QA JSON 文件路径 (顶层 list[dict])",
    )
    parser.add_argument(
        "--aier-qa-path",
        type=str,
        required=True,
        help="Aier 数据过滤后的 Simple QA JSON 文件路径 (顶层 list[dict])",
    )
    parser.add_argument(
        "--out-dir",
        "-o",
        type=str,
        required=True,
        help="输出目录, 将创建 stage1_simple_qa.json / dataset_info.json / stage1_simple_qa_stats.json",
    )
    return parser.parse_args()


def main() -> None:
    """主函数: 读取两个数据集, 合并输出, 生成 dataset_info 与统计信息。"""
    args = parse_args()

    med_path = Path(args.med_qa_path)
    aier_path = Path(args.aier_qa_path)
    out_dir = Path(args.out_dir)

    # 创建输出目录
    try:
        out_dir.mkdir(parents=True, exist_ok=True)
    except Exception as e:
        raise RuntimeError(f"创建输出目录失败: out_dir={out_dir}, error={e}") from e

    # 1) 读取两个 Alpaca 数据集
    med_records = _load_alpaca_json(med_path, source_name="med_qa")
    aier_records = _load_alpaca_json(aier_path, source_name="aier_qa")

    # 2) 合并数据
    combined_records: List[Dict[str, Any]] = []
    combined_records.extend(med_records)
    combined_records.extend(aier_records)

    logger.info(
        "合并完成: med_qa=%d, aier_qa=%d, 总计=%d",
        len(med_records),
        len(aier_records),
        len(combined_records),
    )

    # 3) 写出 stage1_simple_qa.json
    stage1_json_path = out_dir / "stage1_simple_qa.json"
    try:
        stage1_json_path.write_text(
            json.dumps(combined_records, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("合并后的数据已写出: %s", stage1_json_path)
    except Exception as e:
        raise RuntimeError(
            f"写出合并数据 JSON 失败: path={stage1_json_path}, error={e}"
        ) from e

    # 4) 写出 dataset_info.json
    dataset_info = _build_dataset_info()
    dataset_info_path = out_dir / "dataset_info.json"
    try:
        dataset_info_path.write_text(
            json.dumps(dataset_info, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("dataset_info.json 已写出: %s", dataset_info_path)
    except Exception as e:
        raise RuntimeError(
            f"写出 dataset_info.json 失败: path={dataset_info_path}, error={e}"
        ) from e

    # 5) 统计信息并写出 stats JSON
    stats = _compute_stats(
        med_records=med_records,
        aier_records=aier_records,
        combined_records=combined_records,
    )
    stats_path = out_dir / "stage1_simple_qa_stats.json"
    try:
        stats_path.write_text(
            json.dumps(stats, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        logger.info("统计信息已写出: %s", stats_path)
    except Exception as e:
        raise RuntimeError(
            f"写出统计 JSON 失败: path={stats_path}, error={e}"
        ) from e

    # 同时在标准输出上打印一份统计 (方便快速查看)
    print(json.dumps(stats, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
